#pragma once
#ifndef OBR_QUATERNION_H
#define OBR_QUATERNION_H

#include "Matrix4x4f.h"
#include "helper_math.h"
#include <assert.h>

namespace OBR
{
struct Quaternion
{
    float x, y, z, w;

    __forceinline__ __device__ __host__ Quaternion(float _x = 0, float _y = 0, float _z = 0,
                                                   float _w = 1)
        : x(_x), y(_y), z(_z), w(_w)
    {
    }

    __forceinline__ __device__ __host__ bool isValid()
    {
        return (x * x + y * y + z * z + w * w) == 1;
    }

    /**
     * @brief combine two quaternions
     *
     * @param rhs
     */
    __forceinline__ __device__ __host__ Quaternion operator*(Quaternion &rhs)
    {

        return Quaternion(                                 // quaternion mult
            w * rhs.x + x * rhs.w + y * rhs.z - z * rhs.y, // new x
            w * rhs.y - x * rhs.z + y * rhs.w + z * rhs.x, // new y
            w * rhs.z + x * rhs.y - y * rhs.x + z * rhs.w, // new z
            w * rhs.w - x * rhs.x - y * rhs.y - z * rhs.z  // new w
        );
    }

    /**
     * @brief apply quaternion transform
     *
     * @param rhs
     */
    __forceinline__ __device__ __host__ float3 operator*(float3 pos)
    {
        Matrix4x4f mat = getTransformMatrix();
        return mat * pos;
    }

    /**
     * @brief get rotation transform matrix
     *
     * @return __forceinline__
     */
    __forceinline__ __device__ __host__ Matrix4x4f getTransformMatrix()
    {
        assert(this->isValid());
        float c = w;
        float c2 = c * c;
        float s2 = 1 - c2;
        float s = sqrtf(s2);
        float sc = s * c;

        Matrix4x4f mat(1.0f);

        mat.data[0] = make_float4(c2 - (1 - 2 * x * x) * s2, 2 * x * y * s2 + x * z * sc,
                                  2 * x * z * s2 - 2 * y * sc, 0);
        mat.data[1] = make_float4(2 * x * y * s2 - 2 * z * sc, c2 - (1 - 2 * y * y) * s2,
                                  2 * y * z * s2 + 2 * x * sc, 0);
        mat.data[2] = make_float4(2 * x * z * s2 + 2 * y * sc, 2 * y * z * s2 - 2 * x * sc,
                                  c2 - (1 - 2 * z * z) * s2, 0);
        return mat;
    }

    /**
     * @brief get corresponding euler angles
     * @a https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles
     */
    __forceinline__ __device__ __host__ float3 getEuler()
    {
        float3 angles;

        // roll (x-axis rotation)
        double sinr_cosp = 2 * (w * x + y * z);
        double cosr_cosp = 1 - 2 * (x * x + y * y);
        angles.x = std::atan2(sinr_cosp, cosr_cosp);

        // pitch (y-axis rotation)
        double sinp = 2 * (w * y - z * x);
        if (std::abs(sinp) >= 1)
            angles.y = std::copysign(M_PI / 2, sinp); // use 90 degrees if out of range
        else
            angles.y = std::asin(sinp);

        // yaw (z-axis rotation)
        double siny_cosp = 2 * (w * z + x * y);
        double cosy_cosp = 1 - 2 * (y * y + z * z);
        angles.z = std::atan2(siny_cosp, cosy_cosp);

        return angles;
    }

    /**
     * @brief build quaternion from euler angles, following X -> Y -> Z order
     * @a https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles
     *
     * @param euler
     */
    __forceinline__ __device__ __host__ static Quaternion fromEulerAngles(float3 euler)
    {
        /* --------------------------- intuitively -------------------------- */
        // return Quaternion::fromAxisAngle(make_float3(0, 0, 1), euler.z) *
        //        Quaternion::fromAxisAngle(make_float3(0, 1, 0), euler.y) *
        //        Quaternion::fromAxisAngle(make_float3(1, 0, 0), euler.x);

        /* ---------------------------- directly ---------------------------- */
        // Abbreviations for the various angular functions
        double cy = cos(euler.z * 0.5);
        double sy = sin(euler.z * 0.5);
        double cp = cos(euler.y * 0.5);
        double sp = sin(euler.y * 0.5);
        double cr = cos(euler.x * 0.5);
        double sr = sin(euler.x * 0.5);

        Quaternion q;
        q.w = cr * cp * cy + sr * sp * sy;
        q.x = sr * cp * cy - cr * sp * sy;
        q.y = cr * sp * cy + sr * cp * sy;
        q.z = cr * cp * sy - sr * sp * cy;

        return q;
    }

    /**
     * @brief get rotation axis and angle
     *
     * @return (as reference in params) angle, axis
     */
    __forceinline__ __device__ __host__ void getAxisAngle(float3 &axis, float &angle)
    {
        assert(this->isValid());
        float sintheta2 = sqrtf(1 - w * w);
        angle = 2.0f * acosf(w);
        axis = make_float3(x / sintheta2, y / sintheta2, z / sintheta2);
    }

    /**
     * @brief build quaternion from rotation axis and angle
     *
     * @param axis
     * @param angle
     */
    __forceinline__ __device__ __host__ static Quaternion fromAxisAngle(float3 axis, float theta)
    {
        float sintheta2 = sinf(theta / 2);
        float costheta2 = cosf(theta / 2);
        axis = normalize(axis);
        return Quaternion(axis.x * sintheta2, axis.y * sintheta2, axis.z * sintheta2, costheta2);
    }
};
} // namespace OBR

#endif
